import librosa
from pydub import AudioSegment
import cupy as cp
import numpy as np
import torch
import pickle
import os
import threading
import os
import datetime
import warnings
warnings.filterwarnings('ignore')
pathLock=threading.Lock()
basepath=r"D:\\old\\Desktop\\old\\animal\\barking-emotion-recognition\\data\\audioset_audios\\"
printlock=threading.Lock()
denoLock = threading.Lock()
def loadmp3(path):
    audio = AudioSegment.from_file(path)
    if audio.duration_seconds < 1:
        audio += AudioSegment.silent(duration=1000)
    y, sr = librosa.load(audio.export(format='wav'), sr=None, offset=0)
    return y,sr
def denoise(n, y, y_gpu,sr):
    indexs = librosa.effects.split(y, top_db=25 - n)
    noisefrequencies = []

    # 统计频率
    for i in range(len(indexs)):
        start = indexs[i][0]
        end = indexs[i][1]
        noisefrequencies.extend(y_gpu[start:end])
    noisefrequencies=cp.array(noisefrequencies)
    # 移除频率
    denoised_signal = cp.where(~cp.isin(y_gpu, noisefrequencies))[0]

    return cp.asnumpy(denoised_signal).astype('float')
def normalize_audio(x):
    # 将音频信号归一化到 [-1, 1] 范围
    max_amplitude = torch.max(torch.abs(x))
    normalized_x = x / max_amplitude
    return normalized_x
def normalize_audio_cpu(x):
    # 将音频信号归一化到 [-1, 1] 范围
    max_amplitude = cp.max(cp.abs(x))
    normalized_x = x / max_amplitude
    return cp.asnumpy(normalized_x)
def get_wav(path):
    with printlock:
        print(f"start:{path}:{datetime.datetime.now()}")
    with pathLock:
        y, sr = loadmp3(basepath + path + '.mp3')
    y_gpu = cp.array(y)
    with denoLock:
        y_tmp = denoise(10, y, y_gpu, sr)
    try:
        if y_tmp != []:
            y = y_tmp
    except:
        y=y_tmp
    #print(y)
    a = librosa.stft(y, n_fft=16, hop_length=None, win_length=None, window='hann', center=True, pad_mode='reflect')
    b = librosa.amplitude_to_db(abs(a))
    c = librosa.feature.rms(y=y, frame_length=4, hop_length=16)
    a = torch.from_numpy(a)
    b = torch.from_numpy(b)
    c = torch.from_numpy(c)
    a = normalize_audio(a)
    b = normalize_audio(b)
    c = normalize_audio(c)
    with pathLock:
        with open(basepath + r'pickle\\' + path + '.pickle', 'wb') as f:
            pickle.dump((a, b, c), f)
    with printlock:
        print(f"ok:{path}:{datetime.datetime.now()}")
if __name__ == "__main__":
    #get_wav('-7Y8Cmjdvuk_20_30_cut')
    threadlist=[]
    dir_list = os.listdir(basepath)
    for filename in dir_list:
        file = os.path.join(basepath, filename)
        if 'mp3' in file.lower():
            a = threading.Thread(target=get_wav, args=(os.path.splitext(filename)[0],))
            if len(threadlist) > os.cpu_count():
                threadlist[0].join()
                threadlist.pop(0)
            else:
                threadlist.append(a)
            a.start()
    for i in threadlist:
        i.join()
    with printlock:
        
        print(f"success:{datetime.datetime.now()}")